""" data 文件构成
data
  - dataset
    - dataset1
      - train.json
      - val.json
      - test.json
      - buff
  - train_info
    - train_flag
      - model 存储模型文件夹
        - train_flag
      - train.log
"""

import os

cur_dir = os.path.dirname(os.path.realpath(__file__))

project_dir = os.path.normpath(os.path.join(cur_dir, "../.."))

data_dir = os.path.join(project_dir, "data")
dataset_base_dir = os.path.join(data_dir, "dataset")

if "MacBook-Air" in os.popen("hostname").read():
    train_info_dir = os.path.join(data_dir, "train_info")
    model_base_dir = os.path.join(data_dir, "models")
else:
    train_info_dir = "/root/autodl-tmp/train_info"
    model_base_dir = "/root/autodl-tmp/models"


def __check_dir_exists__(dir_list):
    for d in dir_list:
        if not os.path.exists(d):
            os.makedirs(d)


need_check_dir_list = [data_dir, dataset_base_dir, train_info_dir]
__check_dir_exists__(dir_list=need_check_dir_list)


def _makedir_if_not_exists_(file_dir, makedir_if_not_exists=True):
    if makedir_if_not_exists and (not os.path.exists(file_dir)):
        os.makedirs(file_dir)
    return file_dir


def get_dateset_dir(dataset_name, makedir_if_not_exists=False):
    dataset_dir_ = os.path.join(dataset_base_dir, dataset_name)
    return _makedir_if_not_exists_(file_dir=dataset_dir_, makedir_if_not_exists=makedir_if_not_exists)


# ------------------ train info begin ---------------------- #


def get_train_info_dir(train_flag, makedir_if_not_exists=True):
    """ 获取训练信息存储目录 """
    train_info_dir_ = os.path.join(train_info_dir, train_flag)
    return _makedir_if_not_exists_(file_dir=train_info_dir_, makedir_if_not_exists=makedir_if_not_exists)


def get_train_log_path(train_flag, makedir_if_not_exists=True):
    """ 获取日志文件地址 """
    train_info_dir_ = get_train_info_dir(train_flag=train_flag, makedir_if_not_exists=makedir_if_not_exists)
    train_log_path_ = os.path.join(train_info_dir_, "train.log")
    return train_log_path_


def _get_train_model_dir_(train_flag, makedir_if_not_exists=True):
    train_info_dir_ = get_train_info_dir(train_flag=train_flag, makedir_if_not_exists=makedir_if_not_exists)
    save_model_dir_ = os.path.join(train_info_dir_, "model")
    return _makedir_if_not_exists_(file_dir=save_model_dir_, makedir_if_not_exists=makedir_if_not_exists)


def get_train_model_path(train_flag, epoch, makedir_if_not_exists=True):
    """ 获取模型文件地址 """
    save_model_dir_ = _get_train_model_dir_(train_flag=train_flag, makedir_if_not_exists=makedir_if_not_exists)
    save_model_name_ = f"torch_model_{train_flag}_epoch_{epoch:03}.bin"
    return os.path.join(save_model_dir_, save_model_name_)


def get_train_best_model_path(train_flag, makedir_if_not_exists=True):
    save_model_dir_ = _get_train_model_dir_(train_flag=train_flag, makedir_if_not_exists=makedir_if_not_exists)
    save_model_name_ = f"pytorch_model.bin"
    return os.path.join(save_model_dir_, save_model_name_)


def get_train_model_config_path(train_flag, makedir_if_not_exists=True):
    save_model_dir_ = _get_train_model_dir_(train_flag=train_flag, makedir_if_not_exists=makedir_if_not_exists)
    save_model_name_ = f"config.json"
    return os.path.join(save_model_dir_, save_model_name_)


def get_last_save_model_path(train_flag, makedir_if_not_exists=True):
    """ 获取最后存储模型 """
    save_model_dir_ = _get_train_model_dir_(train_flag=train_flag, makedir_if_not_exists=makedir_if_not_exists)
    files_ = os.listdir(save_model_dir_)
    files_ = sorted([f for f in files_ if f.endswith('.bin')])
    if len(files_) <= 0:
        return None
    else:
        last_model_name_ = files_[-1]
        return os.path.join(save_model_dir_, last_model_name_)


def get_last_save_model_epoch_no(train_flag, makedir_if_not_exists=True):
    """ 获取最后存储模型的epoch """
    last_model_path = get_last_save_model_path(train_flag=train_flag, makedir_if_not_exists=makedir_if_not_exists)
    if last_model_path is None:
        return 0
    else:
        try:
            return int(last_model_path.split(".")[0].split('_')[-1])
        except Exception as e:
            return 0

# ------------------ train info end ---------------------- #



